{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "# !pip install yfinance crewai crewai-tools"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [],
   "source": [
    "import re\n",
    "import json\n",
    "import os\n",
    "import yfinance as yf\n",
    "from crewai import Agent, Task, Crew, Process, LLM\n",
    "from crewai_tools import CodeInterpreterTool, FileReadTool"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "from pydantic import BaseModel, Field\n",
    "\n",
    "class QueryAnalysisOutput(BaseModel):\n",
    "    \"\"\"Structured output for the query analysis task.\"\"\"\n",
    "    symbol: str = Field(..., description=\"Stock ticker symbol (e.g., TSLA, AAPL).\")\n",
    "    timeframe: str = Field(..., description=\"Time period (e.g., '1d', '1mo', '1y').\")\n",
    "    action: str = Field(..., description=\"Action to be performed (e.g., 'fetch', 'plot').\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "llm = LLM(\n",
    "    model=\"ollama/deepseek-r1:7b\",\n",
    "    base_url=\"http://localhost:11434\",\n",
    "    # temperature=0.7\n",
    ")\n",
    "\n",
    "# 1) Query parser agent\n",
    "query_parser_agent = Agent(\n",
    "    role=\"Stock Data Analyst\",\n",
    "    goal=\"Extract stock details and fetch required data from this user query: {query}.\",\n",
    "    backstory=\"You are a financial analyst specializing in stock market data retrieval.\",\n",
    "    llm=llm,\n",
    "    verbose=True,\n",
    "    memory=True,\n",
    ")\n",
    "\n",
    "query_parsing_task = Task(\n",
    "    description=\"Analyze the user query and extract stock details.\",\n",
    "    expected_output=\"A dictionary with keys: 'symbol', 'timeframe', 'action'.\",\n",
    "    output_pydantic=QueryAnalysisOutput,\n",
    "    agent=query_parser_agent,\n",
    ")\n",
    "\n",
    "\n",
    "# 2) Code writer agent\n",
    "code_writer_agent = Agent(\n",
    "    role=\"Senior Python Developer\",\n",
    "    goal=\"Write Python code to visualize stock data.\",\n",
    "    backstory=\"\"\"You are a Senior Python developer specializing in stock market data visualization. \n",
    "                 You are also a Pandas, Matplotlib and yfinance library expert.\n",
    "                 You are skilled at writing production-ready Python code\"\"\",\n",
    "    llm=llm,\n",
    "    verbose=True,\n",
    ")\n",
    "\n",
    "code_writer_task = Task(\n",
    "    description=\"\"\"Write Python code to visualize stock data based on the inputs from the stock analyst\n",
    "                   where you would find stock symbol, timeframe and action.\"\"\",\n",
    "    expected_output=\"A clean and executable Python script file (.py) for stock visualization.\",\n",
    "    agent=code_writer_agent,\n",
    ")\n",
    "\n",
    "\n",
    "# 3) Code interpreter agent (uses code interpreter tool from crewai)\n",
    "code_interpreter_tool = CodeInterpreterTool()\n",
    "\n",
    "code_execution_agent = Agent(\n",
    "    role=\"Senior Code Execution Expert\",\n",
    "    goal=\"Review and execute the generated Python code by code writer agent to visualize stock data.\",\n",
    "    backstory=\"You are a code execution expert. You are skilled at executing Python code.\",\n",
    "    # tools=[code_interpreter_tool],\n",
    "    allow_code_execution=True,   # This automatically adds the CodeInterpreterTool\n",
    "    llm=llm,\n",
    "    verbose=True,\n",
    ")\n",
    "\n",
    "code_execution_task = Task(\n",
    "    description=\"\"\"Review and execute the generated Python code by code writer agent to visualize stock data.\"\"\",\n",
    "    expected_output=\"A clean and executable Python script file (.py) for stock visualization.\",\n",
    "    agent=code_execution_agent,\n",
    ")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\u001b[1m\u001b[95m# Agent:\u001b[00m \u001b[1m\u001b[92mStock Data Analyst\u001b[00m\n",
      "\u001b[95m## Task:\u001b[00m \u001b[92mAnalyze the user query and extract stock details.\u001b[00m\n",
      "\n",
      "\n",
      "\u001b[1m\u001b[95m# Agent:\u001b[00m \u001b[1m\u001b[92mStock Data Analyst\u001b[00m\n",
      "\u001b[95m## Final Answer:\u001b[00m \u001b[92m\n",
      "<think>\n",
      "Alright, let me try to figure out what I need to do here. The user has given me a task as a stock data analyst where I have to extract specific details from their query. The query is \"Plot YTD stock gain of Tesla.\" \n",
      "\n",
      "First, I'll break down the query into parts to understand each component. The main goal is to plot (create a graph) of the Year-to-Date (YTD) stock gain for Tesla.\n",
      "\n",
      "Looking at this, there are three key elements that stand out:\n",
      "1. **Symbol**: This should be the ticker symbol for Tesla. From what I remember, Tesla's ticker is TSLA.\n",
      "2. **Timeframe**: The query mentions \"YTD,\" which stands for Year-to-Date, so the timeframe is definitely YTD.\n",
      "3. **Action**: The user wants a plot or graph of the stock gain over this period.\n",
      "\n",
      "Now, considering these elements:\n",
      "- **Symbol**: TSLA\n",
      "- **Timeframe**: YTD\n",
      "- **Action**: Plot (which refers to generating a graph)\n",
      "\n",
      "I need to make sure each part is correctly identified and formatted in the dictionary as per the instructions. I should avoid any markdown formatting and just present it in plain JSON-like structure.\n",
      "\n",
      "So, putting it all together, the final dictionary should look like this:\n",
      "\n",
      "{\n",
      "  \"symbol\": \"TSLA\",\n",
      "  \"timeframe\": \"YTD\",\n",
      "  \"action\": \"Plot\"\n",
      "}\n",
      "\n",
      "I think that covers everything the user is asking for. They wanted the action to be extracting these details and presenting them in a specific format without any additional information or code blocks. So, this should meet their requirements perfectly.\n",
      "</think>\n",
      "\n",
      "{\n",
      "  \"symbol\": \"TSLA\",\n",
      "  \"timeframe\": \"YTD\",\n",
      "  \"action\": \"Plot\"\n",
      "}\u001b[00m\n",
      "\n",
      "\n",
      "\u001b[1m\u001b[95m# Agent:\u001b[00m \u001b[1m\u001b[92mSenior Python Developer\u001b[00m\n",
      "\u001b[95m## Task:\u001b[00m \u001b[92mWrite Python code to visualize stock data based on the inputs from the stock analyst\n",
      "                   where you would find stock symbol, timeframe and action.\u001b[00m\n",
      "\n",
      "\n",
      "\u001b[1m\u001b[95m# Agent:\u001b[00m \u001b[1m\u001b[92mSenior Python Developer\u001b[00m\n",
      "\u001b[95m## Final Answer:\u001b[00m \u001b[92m\n",
      "The complete JSON configuration with the extracted data from the query.\n",
      "\n",
      "```json\n",
      "{\n",
      "  \"symbol\": \"TSLA\",\n",
      "  \"timeframe\": \"YTD\",\n",
      "  \"action\": \"Plot\"\n",
      "}\n",
      "```\u001b[00m\n",
      "\n",
      "\n",
      "\u001b[1m\u001b[95m# Agent:\u001b[00m \u001b[1m\u001b[92mSenior Code Execution Expert\u001b[00m\n",
      "\u001b[95m## Task:\u001b[00m \u001b[92mReview and execute the generated Python code by code writer agent to visualize stock data.\u001b[00m\n",
      "\n",
      "\n",
      "\u001b[1m\u001b[95m# Agent:\u001b[00m \u001b[1m\u001b[92mSenior Code Execution Expert\u001b[00m\n",
      "\u001b[95m## Final Answer:\u001b[00m \u001b[92m\n",
      "```python\n",
      "# This Python code will plot YTD stock data for Tesla using historical prices from Yahoo Finance\n",
      "import yfinance as yf\n",
      "import pandas as pd\n",
      "import matplotlib.pyplot as plt\n",
      "\n",
      "# Extracting the required information from the JSON configuration\n",
      "config = {\n",
      "    \"symbol\": \"TSLA\",\n",
      "    \"timeframe\": \"YTD\",\n",
      "    \"action\": \"Plot\"\n",
      "}\n",
      "\n",
      "# Fetching historical data for Tesla from Yahoo Finance\n",
      "tesla_data = yfinance.download(\n",
      "    tickers=config['symbol'],\n",
      "    period=config['timeframe'],\n",
      "    interval='1d'\n",
      ")\n",
      "\n",
      "# Creating a simple plot of the stock prices over time\n",
      "plt.figure(figsize=(10, 6))\n",
      "plt.plot(tesla_data.index, tesla_data['Adj Close'])\n",
      "plt.title(f'Tesla Stock Price (TSLA) - {config['timeframe']}')\n",
      "plt.xlabel('Date')\n",
      "plt.ylabel('Adjusted Closing Price ($)')\n",
      "plt.grid(True)\n",
      "plt.show()\n",
      "```\u001b[00m\n",
      "\n",
      "\n"
     ]
    }
   ],
   "source": [
    "### --- CREW SETUP --- ###\n",
    "\n",
    "crew = Crew(\n",
    "    agents=[query_parser_agent, code_writer_agent, code_execution_agent],\n",
    "    tasks=[query_parsing_task, code_writer_task, code_execution_task],\n",
    "    process=Process.sequential\n",
    ")\n",
    "\n",
    "# Run the crew with an example query\n",
    "result = crew.kickoff(inputs={\"query\": \"Plot YTD stock gain of Tesla\"})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "```python\n",
      "# This Python code will plot YTD stock data for Tesla using historical prices from Yahoo Finance\n",
      "import yfinance as yf\n",
      "import pandas as pd\n",
      "import matplotlib.pyplot as plt\n",
      "\n",
      "# Extracting the required information from the JSON configuration\n",
      "config = {\n",
      "    \"symbol\": \"TSLA\",\n",
      "    \"timeframe\": \"YTD\",\n",
      "    \"action\": \"Plot\"\n",
      "}\n",
      "\n",
      "# Fetching historical data for Tesla from Yahoo Finance\n",
      "tesla_data = yfinance.download(\n",
      "    tickers=config['symbol'],\n",
      "    period=config['timeframe'],\n",
      "    interval='1d'\n",
      ")\n",
      "\n",
      "# Creating a simple plot of the stock prices over time\n",
      "plt.figure(figsize=(10, 6))\n",
      "plt.plot(tesla_data.index, tesla_data['Adj Close'])\n",
      "plt.title(f'Tesla Stock Price (TSLA) - {config['timeframe']}')\n",
      "plt.xlabel('Date')\n",
      "plt.ylabel('Adjusted Closing Price ($)')\n",
      "plt.grid(True)\n",
      "plt.show()\n",
      "```\n"
     ]
    }
   ],
   "source": [
    "print(result.raw)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": ".venv",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.1"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
